#!/usr/bin/python3
# -*-coding:utf-8 -*-

# Reference:https://github.com/raghakot/keras-resnet.git
# @Time    : 6/28/2019 3:40 PM
# @Author  : Gaopeng.Bai
# @File    : Train.py
# @User    : baigaopeng
# @Software: PyCharm
# Github:https://github.com/Gaopeng-Bai/keras-resnet.git
"""
Adapted from keras example cifar10_cnn.py
Train ResNet-18 on the CIFAR10 small images dataset.

GPU run command with Theano backend (with TensorFlow, the GPU is automatically used):
    THEANO_FLAGS=mode=FAST_RUN,device=gpu,floatX=float32 python Model.py
"""
from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import ReduceLROnPlateau, CSVLogger, EarlyStopping
# Tensorboard
from keras.callbacks import TensorBoard
import argparse
import os
import numpy as np
import resnet
from Preprocessing_cifar10 import preprecessiong

def main():
    # deactivate the warnings for "teh tf library wasn't co to use SSE instructions"
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', default="train", help="train, predict")
    parser.add_argument('--model', default="18", help="18, 34, 50, 101, 152")

    parser.add_argument('--batch_size', default=32)
    parser.add_argument('--nb_classes', default=10)
    parser.add_argument('--nb_epoch', default=200)
    parser.add_argument('--data_augmentation', default=True)
    parser.add_argument('--img_rows', default=32)
    parser.add_argument('--img_cols', default=32)
    parser.add_argument('--img_channels', default=3)

    parser.add_argument('--csv', default=False)

    parser.add_argument('--tensorboard_dir', default='./logs')
    args = parser.parse_args()
    if args.mode == 'train':
        training(args)


def training(args):
    global model, csv_logger

    data_loader = preprecessiong(args)
    X_train, Y_train, X_test, Y_test = data_loader.cifar10()

    if args.model == "18":
        model = resnet.ResnetBuilder.build_resnet_18((args.img_channels, args.img_rows, args.img_cols), args.nb_classes)
    elif args.model == "34":
        model = resnet.ResnetBuilder.build_resnet_34((args.img_channels, args.img_rows, args.img_cols), args.nb_classes)
    elif args.model == "50":
        model = resnet.ResnetBuilder.build_resnet_50((args.img_channels, args.img_rows, args.img_cols), args.nb_classes)
    elif args.model == "101":
        model = resnet.ResnetBuilder.build_resnet_101((args.img_channels, args.img_rows, args.img_cols),
                                                      args.nb_classes)
    elif args.model == "152":
        model = resnet.ResnetBuilder.build_resnet_152((args.ResnetBuilderimg_channels, args.img_rows, args.img_cols),
                                                      args.nb_classes)

    model.compile(loss='categorical_crossentropy',
                  optimizer='adam',
                  metrics=['accuracy'])
    # set leaning rate reducer.
    lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1), cooldown=0, patience=5, min_lr=0.5e-6)
    # set stopper for gradient descent when loss not decrease anymore.
    early_stopper = EarlyStopping(min_delta=0.001, patience=10)
    # save recode information of training step by epoch.
    if args.csv:
        csv_logger = CSVLogger('resnet18_cifar10.csv')

    tbCallBack = TensorBoard(log_dir=args.tensorboard_dir,  # log 目录
                             histogram_freq=0,  # 按照何等频率（epoch）来计算直方图，0为不计算
                             batch_size=args.batch_size,  # 用多大量的数据计算直方图
                             write_graph=True,  # 是否存储网络结构图
                             write_grads=True,  # 是否可视化梯度直方图
                             write_images=True,  # 是否可视化参数
                             embeddings_freq=0,
                             embeddings_layer_names=None,
                             embeddings_metadata=None)

    if not args.data_augmentation:
        print('Not using data augmentation.')
        if args.vsv:
            model.fit(X_train, Y_train,
                      batch_size=args.batch_size,
                      nb_epoch=args.nb_epoch,
                      validation_data=(X_test, Y_test),
                      shuffle=True,
                      callbacks=[lr_reducer, early_stopper, csv_logger, tbCallBack])
        else:
            model.fit(X_train, Y_train,
                      batch_size=args.batch_size,
                      nb_epoch=args.nb_epoch,
                      validation_data=(X_test, Y_test),
                      shuffle=True,
                      callbacks=[lr_reducer, early_stopper, tbCallBack])

    else:
        print('Using real-time data augmentation.')
        # This will do preprocessing and realtime data augmentation:
        datagen = ImageDataGenerator(
            featurewise_center=False,  # set input mean to 0 over the dataset
            samplewise_center=False,  # set each sample mean to 0
            featurewise_std_normalization=False,  # divide inputs by std of the dataset
            samplewise_std_normalization=False,  # divide each input by its std
            zca_whitening=False,  # apply ZCA whitening
            rotation_range=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)  # randomly flip images

        # Compute quantities required for featurewise normalization
        # (std, mean, and principal components if ZCA whitening is applied).
        datagen.fit(X_train)

        # Fit the model on the batches generated by datagen.flow().
        if args.csv:
            model.fit_generator(datagen.flow(X_train, Y_train, batch_size=args.batch_size),
                                steps_per_epoch=X_train.shape[0] // args.batch_size,
                                validation_data=(X_test, Y_test),
                                epochs=args.nb_epoch, verbose=1, max_q_size=100,
                                callbacks=[lr_reducer, early_stopper, csv_logger, tbCallBack])
        else:
            model.fit_generator(datagen.flow(X_train, Y_train, batch_size=args.batch_size),
                                steps_per_epoch=X_train.shape[0] // args.batch_size,
                                validation_data=(X_test, Y_test),
                                epochs=args.nb_epoch, verbose=1, max_queue_size=100,
                                callbacks=[lr_reducer, early_stopper, tbCallBack])


if __name__ == '__main__':
    main()
